#include "DHVideo.h"

extern "C"{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavdevice/avdevice.h>
#include <libavfilter/avfilter.h>
#include <libavfilter/buffersrc.h>
#include <libavfilter/buffersink.h>
#include <libavutil/avutil.h>
#include <libavutil/pixfmt.h>
#include <libswresample/swresample.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
#include <libavutil/hwcontext.h>

#include <libavutil/buffer.h>
#include <libavutil/error.h>
#include <libavutil/hwcontext.h>
#include <libavutil/hwcontext_qsv.h>
#include <libavutil/mem.h>
}

#include <thread>
#include <mutex>
#include <list>
#include <libyuv.h>

typedef struct tagStDecode
{
	AVPixelFormat format;
	SwsContext* sws = nullptr;
	AVPacket* pkt = nullptr;
	AVCodecContext *cd= nullptr;
	AVCodecParserContext *parser;
	const AVCodec *codecDecode = nullptr;
	AVFrame *frame = nullptr, *swframe = nullptr;
	void(*recv)(void* obj, unsigned char* data, int& len, int frame_index);
	void* user_obj =nullptr;

	AVFrame *hwframe = nullptr;
	struct DecodeContext{
		AVBufferRef *hw_device_ctx = nullptr;
	}device_ctx;
	const AVCodecHWConfig *config = nullptr;
}_StDecode;
extern AVPixelFormat g_format;

#pragma region CPU解码
bool Decode(void** handle, AVPacket* pkt)
{
	_StDecode* StDecode = (_StDecode*)*handle;
	if(!StDecode->cd) return false;
	int ret = avcodec_send_packet(StDecode->cd, pkt);
	if (ret < 0) return false;
	while (ret >= 0)
	{
		ret = avcodec_receive_frame(StDecode->cd, StDecode->hwframe==nullptr?
			StDecode->frame:StDecode->hwframe);
		if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
			break;
		else if (ret < 0)
			return false;

		if(StDecode->hwframe && av_hwframe_transfer_data(StDecode->frame, StDecode->hwframe, 0)<0)
			return false;
		if (StDecode->frame != nullptr && StDecode->frame->linesize[0]>0) {
			StDecode->swframe->format = StDecode->format;
			StDecode->swframe->width  = StDecode->frame->width;
			StDecode->swframe->height = StDecode->frame->height;
			if (av_frame_get_buffer(StDecode->swframe, 16)<0)	
				av_frame_get_buffer(StDecode->swframe, 1);
			
			if (StDecode->sws == nullptr){
				StDecode->sws = sws_getContext(StDecode->frame->width, StDecode->frame->height, AV_PIX_FMT_NV12,
					StDecode->swframe->width, StDecode->swframe->height, StDecode->format, SWS_FAST_BILINEAR, NULL, NULL, NULL);
			}
			sws_scale(StDecode->sws,(const uint8_t **)StDecode->frame->data,StDecode->frame->linesize,
				0,StDecode->cd->height,StDecode->swframe->data,StDecode->swframe->linesize);
			StDecode->recv(StDecode->user_obj ,StDecode->swframe->data[0], StDecode->swframe->linesize[0],StDecode->cd->frame_number);

			if(StDecode->hwframe) av_frame_unref(StDecode->hwframe);
			if(StDecode->frame) av_frame_unref(StDecode->frame);
			if(StDecode->swframe) av_frame_unref(StDecode->swframe);
		}
	}
	return true;
}
void SWBeginDecode(void** handle, void* obj, int code, void(*decode)(void* obj, unsigned char* data, int& len, int frame_index))
{
	_StDecode* StDecode = (_StDecode*)*handle;
	StDecode->user_obj = obj;
	StDecode->recv = decode;
	StDecode->format = g_format;
	
	StDecode->codecDecode = (code==AV_CODEC_ID_H264?avcodec_find_decoder(AV_CODEC_ID_H264):
		(code==AV_CODEC_ID_H265?avcodec_find_decoder(AV_CODEC_ID_H265):
		avcodec_find_decoder((AVCodecID)code)));
    if(!StDecode->codecDecode) return;
    StDecode->parser = av_parser_init(StDecode->codecDecode->id);
    if(!StDecode->parser) return;
    StDecode->cd = avcodec_alloc_context3(StDecode->codecDecode);
    if(!StDecode->cd) return;
    if (avcodec_open2(StDecode->cd, StDecode->codecDecode, NULL) < 0)
        return;

	StDecode->frame = av_frame_alloc();
	StDecode->swframe = av_frame_alloc();
	StDecode->pkt = av_packet_alloc();
}
void SWImageDecode(void** handle, unsigned char* data, int len)
{
	_StDecode* StDecode = (_StDecode*)*handle;
	if(!StDecode->pkt || !StDecode->parser || !StDecode->cd) return;
    while (len > 0)
    {
        int ret = av_parser_parse2(StDecode->parser, StDecode->cd, &StDecode->pkt->data, 
			&StDecode->pkt->size,data, len, AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
        if(ret<0)   break;
        data += ret;
        len -= ret;
		if(StDecode->pkt->size>0 && !Decode(handle,StDecode->pkt))
			break;
		if(StDecode->pkt) 
            av_packet_unref(StDecode->pkt);
    }
}
void SWEndDecode(void** handle)
{
	Decode(handle,nullptr);
	_StDecode* StDecode = (_StDecode*)*handle;
	if(StDecode->frame) av_frame_free(&StDecode->frame);
	if(StDecode->swframe) av_frame_free(&StDecode->swframe);
	if(StDecode->pkt) av_free_packet(StDecode->pkt);
	if(StDecode->parser) av_parser_close(StDecode->parser);
	if(StDecode->cd) avcodec_close(StDecode->cd);
	if(StDecode->sws) sws_freeContext(StDecode->sws);
	if(StDecode) {
		delete StDecode;
		StDecode =  nullptr;
	}
}
#pragma endregion

#pragma region QSV解码
void QSVBeginDecode(void** handle, void* obj, int code, void(*decode)(void* obj, unsigned char* data, int& len, int frame_index))
{
	_StDecode* StDecode = (_StDecode*)*handle;
	StDecode->user_obj = obj;
	StDecode->recv = decode;
	StDecode->format = g_format;
	
    StDecode->codecDecode = (code==AV_CODEC_ID_H264?avcodec_find_decoder_by_name("h264_qsv"):
		(code==AV_CODEC_ID_H265?avcodec_find_decoder_by_name("hevc_qsv"):
		avcodec_find_decoder((AVCodecID)code)));
    if(!StDecode->codecDecode) return;
    StDecode->parser = av_parser_init(StDecode->codecDecode->id);
    if(!StDecode->parser) return;
    StDecode->cd = avcodec_alloc_context3(StDecode->codecDecode);
    if(!StDecode->cd) return;

	StDecode->cd->opaque = &StDecode->device_ctx;
	StDecode->cd->get_format = [](AVCodecContext *avctx, const enum AVPixelFormat *pix_fmts)->AVPixelFormat{
		while (*pix_fmts != AV_PIX_FMT_NONE) {
			if (*pix_fmts == AV_PIX_FMT_QSV) {
				_StDecode::DecodeContext* decode = (_StDecode::DecodeContext*)avctx->opaque;
				AVHWFramesContext  *frames_ctx;
				AVQSVFramesContext *frames_hwctx;

				/* create a pool of surfaces to be used by the decoder */
				avctx->hw_frames_ctx = av_hwframe_ctx_alloc(decode->hw_device_ctx);
				if (!avctx->hw_frames_ctx)
					return AV_PIX_FMT_NONE;
				frames_ctx   = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
				frames_hwctx = (AVQSVFramesContext*)frames_ctx->hwctx;
				frames_ctx->format            = AV_PIX_FMT_QSV;
				frames_ctx->sw_format         = avctx->sw_pix_fmt;
				frames_ctx->width             = FFALIGN(avctx->coded_width,  32);
				frames_ctx->height            = FFALIGN(avctx->coded_height, 32);
				frames_ctx->initial_pool_size = 32*2;  //4的倍数
				frames_hwctx->frame_type = MFX_MEMTYPE_VIDEO_MEMORY_DECODER_TARGET;
				if (av_hwframe_ctx_init(avctx->hw_frames_ctx) < 0)
					return AV_PIX_FMT_NONE;
				return AV_PIX_FMT_QSV;
			}
			pix_fmts++;
		}
		return AV_PIX_FMT_NONE;
	};
	
    if (avcodec_open2(StDecode->cd, StDecode->codecDecode, NULL) < 0)
        return;

	StDecode->hwframe = av_frame_alloc();
	StDecode->frame = av_frame_alloc();
	StDecode->swframe = av_frame_alloc();
	StDecode->pkt = av_packet_alloc();
}
void QSVImageDecode(void** handle, unsigned char* data, int len)
{
	_StDecode* StDecode = (_StDecode*)*handle;
	if(!StDecode->pkt || !StDecode->parser || !StDecode->cd) return;
	while (len > 0)
    {
        int ret = av_parser_parse2(StDecode->parser, StDecode->cd, &StDecode->pkt->data, 
			&StDecode->pkt->size,data, len, AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
        if(ret<0)   break;
        data += ret;
        len -= ret;
		if(StDecode->pkt->size>0 && !Decode(handle,StDecode->pkt))
			break;

		if(StDecode->pkt) 
            av_packet_unref(StDecode->pkt);
    }
}
void QSVEndDecode(void** handle)
{
	Decode(handle,nullptr);
	_StDecode* StDecode = (_StDecode*)*handle;
	if(StDecode->hwframe) av_frame_free(&StDecode->hwframe);
	if(StDecode->frame) av_frame_free(&StDecode->frame);
	if(StDecode->swframe) av_frame_free(&StDecode->swframe);
	if(StDecode->pkt) av_free_packet(StDecode->pkt);
	if(StDecode->parser) av_parser_close(StDecode->parser);
	if(StDecode->cd) avcodec_close(StDecode->cd);
	if(StDecode->sws) sws_freeContext(StDecode->sws);
	if(StDecode->device_ctx.hw_device_ctx) av_buffer_unref(&StDecode->device_ctx.hw_device_ctx);
	if(StDecode) {
		delete StDecode;
		StDecode =  nullptr;
	}
}
#pragma endregion

#pragma region CUDA解码
AVPixelFormat hw_format;
void CUDABeginDecode(void** handle, void* obj, int code, void(*decode)(void* obj, unsigned char* data, int& len, int frame_index))
{
	_StDecode* StDecode = (_StDecode*)*handle;
	StDecode->user_obj = obj;
	StDecode->recv = decode;
	StDecode->format = g_format;
	
	StDecode->codecDecode = (code==AV_CODEC_ID_H264?avcodec_find_decoder_by_name("h264_cuvid"):
    	(code==AV_CODEC_ID_H265?avcodec_find_decoder_by_name("hevc_cuvid"):
		avcodec_find_decoder((AVCodecID)code)));
    if(!StDecode->codecDecode) return;
    StDecode->parser = av_parser_init(StDecode->codecDecode->id);
    if(!StDecode->parser) return;
    StDecode->cd = avcodec_alloc_context3(StDecode->codecDecode);
    if(!StDecode->cd) return;
	
	for (int i = 0;; i++) {
		StDecode->config = avcodec_get_hw_config(StDecode->codecDecode, i);
		if (!StDecode->config) break;
		if(StDecode->config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) {
			if(StDecode->config->device_type == AV_HWDEVICE_TYPE_CUDA){
				hw_format = StDecode->config->pix_fmt;
				break;
			}
			else{
				StDecode->config = nullptr;
			}
		}
	}
	StDecode->cd->get_format = [](AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts)->AVPixelFormat{
		const enum AVPixelFormat *p;
		for (p = pix_fmts; *p != -1; p++) {
			if (*p == hw_format)
				return *p;
		}
		return AV_PIX_FMT_NONE;
	};
	StDecode->cd->hw_device_ctx = av_buffer_ref(StDecode->device_ctx.hw_device_ctx);
    if (avcodec_open2(StDecode->cd, StDecode->codecDecode, NULL) < 0)
        return;

	StDecode->hwframe = av_frame_alloc();
	StDecode->frame = av_frame_alloc();
	StDecode->swframe = av_frame_alloc();
	StDecode->pkt = av_packet_alloc();
}
void CUDAImageDecode(void** handle, unsigned char* data, int len)
{
	_StDecode* StDecode = (_StDecode*)*handle;
	if(!StDecode->pkt || !StDecode->parser || !StDecode->cd) return;
	while (len > 0)
    {
        int ret = av_parser_parse2(StDecode->parser, StDecode->cd, &StDecode->pkt->data, 
			&StDecode->pkt->size,data, len, AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
        if(ret<0)   break;
        data += ret;
        len -= ret;
		if(StDecode->pkt->size>0 && !Decode(handle, StDecode->pkt))
			break;
		
		if(StDecode->pkt) 
            av_packet_unref(StDecode->pkt);
    }
}
void CUDAEndDecode(void** handle)
{
	Decode(handle,nullptr);
	_StDecode* StDecode = (_StDecode*)*handle;
	if(StDecode->hwframe) av_frame_free(&StDecode->hwframe);
	if(StDecode->frame) av_frame_free(&StDecode->frame);
	if(StDecode->swframe) av_frame_free(&StDecode->swframe);
	if(StDecode->pkt) av_free_packet(StDecode->pkt);
	if(StDecode->parser) av_parser_close(StDecode->parser);
	if(StDecode->cd) avcodec_close(StDecode->cd);
	if(StDecode->sws) sws_freeContext(StDecode->sws);
	if(StDecode->device_ctx.hw_device_ctx) av_buffer_unref(&StDecode->device_ctx.hw_device_ctx);
	if(StDecode) {
		delete StDecode;
		StDecode =  nullptr;
	}
}
#pragma endregion

extern int n_hw_accels;
extern AVHWDeviceType hwtype;
void BeginDecode(void** handle, void* obj, int code, void(*decode)(void* obj, unsigned char* data, int& len, int frame_index))
{
	_StDecode* StDecode = new _StDecode();
    *handle = (void*)StDecode;
	av_log_set_level(AV_LOG_ERROR);
	if(n_hw_accels>0){
        if(n_hw_accels == 1){
            hwtype = AV_HWDEVICE_TYPE_NONE;
            SWBeginDecode(handle,obj,code,decode);
        }
        else if(n_hw_accels == 2){
            hwtype = AV_HWDEVICE_TYPE_QSV;
            if(!av_hwdevice_ctx_create(&StDecode->device_ctx.hw_device_ctx, AV_HWDEVICE_TYPE_QSV, "auto", NULL, 0))
                QSVBeginDecode(handle,obj,code,decode);
        }
        else if(n_hw_accels == 3){
            hwtype = AV_HWDEVICE_TYPE_CUDA;
            if(!av_hwdevice_ctx_create(&StDecode->device_ctx.hw_device_ctx, AV_HWDEVICE_TYPE_CUDA, "auto", NULL, 0))
                CUDABeginDecode(handle,obj,code,decode);
        }
    }
	else{
		if((hwtype = av_hwdevice_find_type_by_name("cuda")) == AV_HWDEVICE_TYPE_CUDA 
			&&!av_hwdevice_ctx_create(&StDecode->device_ctx.hw_device_ctx, AV_HWDEVICE_TYPE_CUDA, "auto", NULL, 0))
			CUDABeginDecode(handle,obj,code,decode);
		else if((hwtype = av_hwdevice_find_type_by_name("qsv")) == AV_HWDEVICE_TYPE_QSV 
			&&!av_hwdevice_ctx_create(&StDecode->device_ctx.hw_device_ctx, AV_HWDEVICE_TYPE_QSV, "auto", NULL, 0))
			QSVBeginDecode(handle,obj,code,decode);
		else if((hwtype = AV_HWDEVICE_TYPE_NONE) == AV_HWDEVICE_TYPE_NONE)
			SWBeginDecode(handle,obj,code,decode);
	}
}
void ImageDecode(void** handle, unsigned char* data, int len)
{
    if(hwtype == AV_HWDEVICE_TYPE_CUDA)
		CUDAImageDecode(handle,data,len);
	else if(hwtype == AV_HWDEVICE_TYPE_QSV)
		QSVImageDecode(handle,data,len);
    else
		SWImageDecode(handle,data,len);
}
void EndDecode(void** handle)
{
    if(hwtype == AV_HWDEVICE_TYPE_CUDA)
		CUDAEndDecode(handle);
	else if(hwtype == AV_HWDEVICE_TYPE_QSV)
		QSVEndDecode(handle);
    else
		SWEndDecode(handle);
}